/*
 *        (C) COPYRIGHT Ingenic Limited.
 *             ALL RIGHTS RESERVED
 *
 * File       : inference_yolov5s.cpp
 * Authors    : ljhong, cywang
 * Create Time: 2024-11-21:15:57:21
 * Description:
 * 
 */

 #include "ingenic_mnni.h"
 #include "common/common_utils.h"
 #include "utils/imgproc.h"
 #include "utils/postproc.h"
 
 #include <cmath>
 #include <errno.h>
 #include <fstream>
 #include <iostream>
 #include <memory>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string>
 #include <cstring>
 #include <string.h>
 #include <vector>
 #include <algorithm>
 #include <sys/stat.h>   
 #include <sys/types.h>
 #include <dirent.h>     
 #include <unistd.h>     
 #include <sstream>      
 #include <iomanip>      
 #include <numeric>
 #include "yolo_interface.h"
 #include "recog_interface.h"
 #include "xb_image_utils.h"

 #ifdef URANUS_NV12
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::NV12;
#else
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::RGBA;
#endif

using namespace Recog;
using namespace XBUtils;

#define RESIZE_FLAG 1

 // 辅助函数：从完整路径中提取文件名
 std::string get_filename_from_path(const std::string& path) {
     size_t last_slash = path.find_last_of("/\\");
     if (last_slash != std::string::npos) {
         return path.substr(last_slash + 1);
     }
     return path;
 }
 
 
bool is_number(const string& s) {
    return !s.empty() && all_of(s.begin(), s.end(), ::isdigit);
}

vector<string> get_image_paths(const string& dir_path) {
    vector<string> image_paths;
    DIR* dir = opendir(dir_path.c_str());
    if (!dir) {
        perror("opendir failed");
        return image_paths;
    }
    
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        string filename = entry->d_name;
        if (filename == "." || filename == "..") continue;
        
        if (filename.find(".jpg") != string::npos || 
            filename.find(".png") != string::npos ||
            filename.find(".jpeg") != string::npos) {
            image_paths.push_back(dir_path + "/" + filename);
        }
    }
    closedir(dir);
    
    // 按数字名称排序
    sort(image_paths.begin(), image_paths.end(), [](const string& a, const string& b) {
        string base_a = a.substr(a.find_last_of('/') + 1);
        string num_a = base_a.substr(0, base_a.find_last_of('.'));
        
        string base_b = b.substr(b.find_last_of('/') + 1);
        string num_b = base_b.substr(0, base_b.find_last_of('.'));
        
        // 如果都是纯数字文件名
        if (is_number(num_a) && is_number(num_b)) {
            return atoi(num_a.c_str()) < atoi(num_b.c_str());
        }
        // 否则按字符串排序
        return base_a < base_b;
    });
    
    return image_paths;
}
 

int remove_directory(const char *path) {
    DIR *dir = opendir(path);
    if (!dir) return -1;
    
    struct dirent *entry;
    while ((entry = readdir(dir)) != NULL) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
            continue;
            
        char full_path[PATH_MAX];
        snprintf(full_path, sizeof(full_path), "%s/%s", path, entry->d_name);
        
        if (entry->d_type == DT_DIR) {
            remove_directory(full_path);
        } else {
            unlink(full_path);
        }
    }
    closedir(dir);
    return rmdir(path);
}

// 辅助函数：创建文件夹（递归）
bool create_directory(const string& path) {
    struct stat st;
    if (stat(path.c_str(), &st) == 0) {
        return S_ISDIR(st.st_mode); // 已存在且是目录
    }
    // 递归创建父目录
    size_t pos = path.find_last_of('/');
    if (pos != string::npos) {
        string parent = path.substr(0, pos);
        if (!create_directory(parent)) return false;
    }
    return mkdir(path.c_str(), 0755) == 0;
}


 int main(int argc, char **argv) {
     if (argc != 5) {
         printf("Usage: %s recog_model_path register_image_path image_folder_path output_dir\n", argv[0]);
         exit(0);
     }
 
     std::string recog_model_path = argv[1];
     std::string register_image_path = argv[2];
     std::string image_folder_path = argv[3];
     std::string output_dir = argv[4];
 
    // 创建输出目录（如果存在则先删除）
    if (access(output_dir.c_str(), F_OK) == 0) {
        if (remove_directory(output_dir.c_str()) != 0) {
            perror("remove existing directory failed");
            return -1;
        }
    }

    if (!create_directory(output_dir)) {
        perror("create output directory failed");
        return -1;
    }

    string input_dir = output_dir + "/input";
    if (!create_directory(input_dir)) {
        perror("create input_dir directory failed");
        return -1;
    }

    string register_dir = output_dir + "/register";
    if (!create_directory(register_dir)) {
        perror("create register_dir directory failed");
        return -1;
    }

     // 创建输出目录
     mkdir(output_dir.c_str(), 0755);
 
     // 打开结果文件
     std::ofstream result_file(output_dir + "/results.txt");
     result_file << "注册图片路径: " << register_image_path << "\n";
     result_file << "测试图片目录: " << image_folder_path << "\n\n";
 
     // 加载识别模型
     RecogConfig recog_config;
     recog_config.dist_thres = 1.0;
     recog_config.model_path = recog_model_path;     
     recog_config.input_w = 224;    
     recog_config.input_h = 224; 
     recog_config.size = 128;
 
     auto recog_interface = RecogInterface(recog_config);
     recog_interface.loadModel();
     recog_interface.print_model_forward_memory();

     FrameInfo input_info;
     MyImg input_img, resize_img, tmp_img, register_img;
    // 读取图片
     ImageUtils::read_img(input_info, input_img, register_image_path);
     magik::uranus::shape_t shape = {1, recog_config.input_h, recog_config.input_w, 4};
     magik::uranus::Tensor recog_working_img(shape);
     string filename = register_image_path.substr(register_image_path.find_last_of('/') + 1);
     string register_path = register_dir + "/" + filename;
     #if RESIZE_FLAG // resize
     ImageUtils::resize_image(input_img, resize_img, recog_config.input_w, recog_config.input_h);
     ImageUtils::img_to_tensor(resize_img, recog_working_img, recog_config.input_w, recog_config.input_h);
     result_file << "图片预处理方法: " << "Resize" << "\n\n";
     #else // letterbox
     result_file << "图片预处理方法: " << "Letterbox" << "\n\n";
     magik::uranus::Box_t tmp_box;
     tmp_box = {0, 0, static_cast<float>(input_info.width), static_cast<float>(input_info.height)};
     ImageUtils::prepare_data(input_info, recog_working_img, 
        tmp_box, 
        recog_config.input_w, 
        recog_config.input_h,
        DATA_FORMAT);        
    #endif   
     RegisterInfo register_info;
     std::vector<RegisterInfo> register_infos;
     recog_interface.registerCat(recog_working_img, register_info, "cat", "123");
     register_infos.push_back(register_info);
     recog_interface.printRegisterInfo(register_infos);

     ImageUtils::tensor_to_img(recog_working_img, register_img);
    ImageUtils::saveImg(register_img, register_path, 60, std::string());
    if(register_img.data != nullptr) ImageUtils::free_stb_mem(register_img.data);


     std::vector<std::vector<float>> feat_vecs;
     vector<string> image_paths = get_image_paths(image_folder_path);
     int same_cat_count = 0;
     int total_images = image_paths.size();
 
     // 创建imgs目录
     mkdir((output_dir + "/imgs").c_str(), 0755);

     for (size_t i = 0; i < image_paths.size(); i++) {
         const auto& image_path = image_paths[i];
         // 创建图片保存路径
         std::string filename = get_filename_from_path(image_path);
         std::string img_save_path = output_dir + "/imgs/" + filename;
         
         // 复制图片到输出目录
         std::string copy_cmd = "cp '" + image_path + "' '" + img_save_path + "'";
         system(copy_cmd.c_str());
 
         // 特征提取
         ImageUtils::read_img(input_info, input_img, image_path);
         int img_width = input_info.width;
         int img_height = input_info.height;
         float aspect_ratio = static_cast<float>(img_width) / img_height;
         if (img_width * img_height < 100 * 100 || aspect_ratio < 0.9f || aspect_ratio > 1.1111f) {
            printf("image %s is not valid, skip\n", image_path.c_str());
            continue;
         }
         #if RESIZE_FLAG // resize
         ImageUtils::resize_image(input_img, resize_img, recog_config.input_w, recog_config.input_h);
         ImageUtils::img_to_tensor(resize_img, recog_working_img, recog_config.input_w, recog_config.input_h);
         #else // letterbox
         magik::uranus::Box_t tmp_box;
         tmp_box = {0, 0, static_cast<float>(input_info.width), static_cast<float>(input_info.height)};
         ImageUtils::prepare_data(input_info, recog_working_img, 
            tmp_box, 
            recog_config.input_w, 
            recog_config.input_h,
            DATA_FORMAT);        
        #endif   
         filename = image_path.substr(image_path.find_last_of('/') + 1);
         string output_path = input_dir + "/" + filename;
         ImageUtils::tensor_to_img(recog_working_img, tmp_img);
         ImageUtils::saveImg(tmp_img, output_path.c_str(), 60, std::string());

         std::vector<RecognizeInfo> recognize_results;
         recog_interface.recognizeCat(recog_working_img, recognize_results, register_infos);
        //  recog_interface.printRecognizeResults(recognize_results);
         result_file << "图片路径: " << image_path << "\n";
         for(auto& result : recognize_results) {
            result_file << "id: " << result.id << ", name: " << result.name << ", distance: " << result.distance << ", is_match: " << result.is_match << "\n\n";
            if(result.is_match) {
                same_cat_count++;
            }
        }
     }
 
     // 计算并保存识别率
     float recognition_rate = static_cast<float>(same_cat_count) / total_images * 100;
     result_file << "\n识别率统计:\n";
     result_file << "总测试图片数: " << total_images << "\n";
     result_file << "正确识别数: " << same_cat_count << "\n";
     result_file << "识别率: " << recognition_rate << "%\n";
 
     std::cout << "\n识别率统计:\n";
     std::cout << "总测试图片数: " << total_images << "\n";
     std::cout << "正确识别数: " << same_cat_count << "\n";
     std::cout << "识别率: " << recognition_rate << "%\n";
    if(resize_img.data != nullptr) ImageUtils::free_stb_mem(resize_img.data);
    if(input_info.data != nullptr)ImageUtils::free_stb_mem(input_info.data);
    if(tmp_img.data != nullptr) ImageUtils::free_stb_mem(tmp_img.data);
        
    recog_interface.releaseModel();
    return 0;
 }